{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Weakly superivised instance segmentation using TAO Mask Auto-labeler\n",
    "\n",
    "[Mask Auto-labeler (MAL)](https://arxiv.org/abs/2301.03992) is a high-quality Transformer-based mask auto-labeling framework for instance segmentation using only box annotations. MAL takes\n",
    "box-cropped images as inputs and conditionally generates their mask pseudo-labels.\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample prediction of MAL model\n",
    "<img align=\"center\" src=\"https://github.com/vpraveen-nv/model_card_images/blob/main/cv/notebook/common/mal_sample.jpg?raw=true\" width=\"960\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Take a pretrained model and train a MAL model on COCO dataset\n",
    "* Evaluate the trained model\n",
    "* Run inference with the trained model and visualize the result\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of MAL using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO launcher](#head-1)\n",
    "2. [Prepare dataset and download pretrained model](#head-2)\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate a trained model](#head-5)\n",
    "6. [Run inference](#head-6)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "When using the purpose-built pretrained models from NGC, please make sure to set the `$KEY` environment variable to the key as mentioned in the model overview. Failing to do so, can lead to errors when trying to load them as pretrained models.\n",
    "\n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users workspace. Please note that the dataset to run this notebook is expected to reside in the `$LOCAL_PROJECT_DIR/data`, while the TAO experiment generated collaterals will be output to `$LOCAL_PROJECT_DIR/mal/`. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.\n",
    "\n",
    "The TAO launcher uses docker containers under the hood, and **for our data and results directory to be visible to the docker, they need to be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options like the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. Here, we can map directories in which we save the data, specs, results and cache. You should configure it for your specific case so these directories are correctly visible to the docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=/path/to/local/tao-experiments\n",
    "\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\")\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"mal\")\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/mal\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p $HOST_DATA_DIR\n",
    "! mkdir -p $HOST_SPECS_DIR\n",
    "! mkdir -p $HOST_RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tlt_configs = {\n",
    "   \"Mounts\":[\n",
    "         # Mapping the Local project directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "       }\n",
    "   ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"16G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         },\n",
    "        # \"user\": \"{}:{}\".format(os.getuid(), os.getgid()),\n",
    "        \"network\": \"host\"\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tlt_configs, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in the `nvidia-pyindex` python index. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python >=3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the virtualenv and virtualenvwrapper packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-pyindex\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset and download pretrained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Prepare dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " We will be using the COCO dataset for the tutorial. The following script will download COCO dataset automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create local dir\n",
    "!mkdir -p $HOST_DATA_DIR\n",
    "# Download the data\n",
    "!bash $HOST_SPECS_DIR/download_coco.sh $HOST_DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verification\n",
    "!ls -l $HOST_DATA_DIR/raw-data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Download pretrained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $LOCAL_PROJECT_DIR/ngccli\n",
    "\n",
    "# Remove any previously existing CLI installations\n",
    "!rm -rf $LOCAL_PROJECT_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $LOCAL_PROJECT_DIR/ngccli\n",
    "!unzip -u \"$LOCAL_PROJECT_DIR/ngccli/$CLI\" -d $LOCAL_PROJECT_DIR/ngccli/\n",
    "!rm -f $LOCAL_PROJECT_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"LOCAL_PROJECT_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List available pretrained models\n",
    "!ngc registry model list nvidia/tao/pretrained_mask_auto_label:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pull pretrained model from NGC\n",
    "!ngc registry model download-version nvidia/tao/pretrained_mask_auto_label:vit-base --dest $LOCAL_PROJECT_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check that model is downloaded into dir.\")\n",
    "!ls -l $LOCAL_PROJECT_DIR/pretrained_mask_auto_label_vvit-base"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide experiment spec file <a class=\"anchor\" id=\"head-3\"></a>\n",
    "\n",
    "We provide a specification file to configure the key parameters for this demo including:\n",
    "\n",
    "* experiment config: configure the global experiment settings\n",
    "    * num_nodes: number of nodes (num_nodes=1 for single node)\n",
    "    * results_dir: the directory where your checkpoints will be saved\n",
    "    * checkpoint: pretrained weights (can be either a pretrained backbone model or a trained MAL model)\n",
    "* dataset config: configure the training and validation datasets\n",
    "    * train_img_dir: annotation file for train data. required to be in COCO json format\n",
    "    * train_ann_path: the root directory for train images\n",
    "    * val_img_dir: the root directory for validation images\n",
    "    * val_ann_path: annotation file for validation data. required to be in COCO json format\n",
    "* model config: configure the model setting\n",
    "    * arch: the backbone architecture for MAL\n",
    "* train_config: configure the training hyperparameters\n",
    "    * lr: learning rate for training the model\n",
    "    * strategy: strategy for distributed training\n",
    "    * batch_size: batch size per gpu\n",
    "    * use_amp: whether to use AMP\n",
    "    * max_epochs: number of epochs\n",
    "    * crop_size: input bounding box size\n",
    "\n",
    "* **Note that the sample spec is not meant to produce SOTA accuracy on COCO. To reproduce SOTA, you might want to use TAO to train an ImageNet model first and follow the original parameters for COCO.**\n",
    "\n",
    "Please refer to the TAO documentation about MAL to get all the parameters that are configurable.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/spec.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* WARNING: COCO training takes about 40+ hours to complete using 8 V100 gpus. As a result, **we highly recommend that you run training with multiple high-end gpus (e.g. V100, A100)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: The following paths are set from the perspective of the TAO Docker.\n",
    "\n",
    "# The data is saved here\n",
    "%env DATA_DIR=/data\n",
    "%env SPECS_DIR=/specs\n",
    "%env RESULTS_DIR=/results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"For multi-GPU, change train.num_gpus in spec.yaml based on your machine.\")\n",
    "!tao model mal train -e $SPECS_DIR/spec.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Model checkpoints:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "# %env NUM_EPOCH=029\n",
    "\n",
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "# tmp=!ls $HOST_RESULTS_DIR/train/*.pth | grep epoch_$NUM_EPOCH\n",
    "# %env CHECKPOINT={tmp[0]}\n",
    "\n",
    "# Or get the latest checkpoint\n",
    "os.environ[\"CHECKPOINT\"] = os.path.join(os.getenv(\"HOST_RESULTS_DIR\"), \"train/mal_model_latest.pth\")\n",
    "\n",
    "print('Rename a trained model: ')\n",
    "print('---------------------')\n",
    "!cp $CHECKPOINT $HOST_RESULTS_DIR/train/mal_model.ckpt\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train/mal_model.ckpt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate a trained model <a class=\"anchor\" id=\"head-5\"></a>\n",
    "\n",
    "In this section, we run the `evaluate` tool to evaluate the trained model and produce the mIOU metric.\n",
    "\n",
    "In `spec.yaml`, we specify a few key parameters for evaluation including:\n",
    "* experiment config\n",
    "    * gpu_ids: gpu indices to use\n",
    "    * checkpoint: a trained MAL model\n",
    "* model config: configure the model setting\n",
    "    * arch: the backbone architecture for MAL\n",
    "* dataset config: configure the training and validation datasets\n",
    "    * val_img_dir: the root directory for validation images\n",
    "    * val_ann_path: annotation file for validation data. required to be in COCO json format\n",
    "* model config: configure the model setting\n",
    "    * arch: the backbone architecture for MAL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate on TAO model\n",
    "!tao model mal evaluate -e $SPECS_DIR/spec.yaml evaluate.checkpoint=$RESULTS_DIR/train/mal_model.ckpt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Run Inference <a class=\"anchor\" id=\"head-6\"></a>\n",
    "In this section, we run the `inference` tool to generate inferences on the trained models and visualize the results. The `inference` tool produces an output annotation json file with pseudo-mask info.\n",
    "\n",
    "In `spec.yaml`, we specify a few key parameters for inference including:\n",
    "* experiment config\n",
    "    * gpu_ids: gpu indices to use\n",
    "    * checkpoint: a trained MAL model\n",
    "* model config: configure the model setting\n",
    "    * arch: the backbone architecture for MAL\n",
    "* dataset config: configure the training and validation datasets\n",
    "    * val_img_dir: the root directory for validation images\n",
    "    * val_ann_path: annotation file for validation data. required to be in COCO json format\n",
    "* model config: configure the model setting\n",
    "    * arch: the backbone architecture for MAL\n",
    "* inference config: configure the data and output for inference\n",
    "    * img_dir: the root directory for test images\n",
    "    * ann_path: annotation file for test data. required to be in COCO json format\n",
    "    * label_dump_path: the output json file with pseudo-mask info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model mal inference -e $SPECS_DIR/spec.yaml inference.checkpoint=$RESULTS_DIR/train/mal_model.ckpt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.1. Visualize the result <a class=\"anchor\" id=\"head-6-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# install deps\n",
    "!pip3 install Cython==0.29.36\n",
    "!pip3 install numpy\n",
    "!pip3 install pillow\n",
    "!pip3 install \"matplotlib>=3.3.3, <4.0\"\n",
    "!pip3 install pycocotools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "matplotlib.use('Agg')\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "from pycocotools.coco import COCO\n",
    "%matplotlib inline\n",
    "\n",
    "img_dir = os.path.join(os.environ['HOST_DATA_DIR'], 'raw-data/val2017/')\n",
    "json_path = os.path.join(os.environ['HOST_RESULTS_DIR'], 'instances_val2017_mal.json')\n",
    "data = COCO(annotation_file=json_path)\n",
    "cat_ids = data.getCatIds()\n",
    "query_id = cat_ids[0] # pick the 1st category\n",
    "# Get the image ids containing the object of the category.\n",
    "img_ids = data.getImgIds(catIds=[query_id])\n",
    "# Pick 1st image\n",
    "img_id = img_ids[0]\n",
    "img_info = data.loadImgs([img_id])[0]\n",
    "img_file_name = img_info[\"file_name\"]\n",
    "print(img_file_name)\n",
    "ann_ids = data.getAnnIds(imgIds=[img_id], iscrowd=None)\n",
    "anns = data.loadAnns(ann_ids)\n",
    "plt.clf()\n",
    "im = Image.open(os.path.join(img_dir, img_file_name))\n",
    "plt.axis(\"off\")\n",
    "plt.imshow(np.asarray(im))\n",
    "data.showAnns(anns, draw_bbox=True)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook has come to an end."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10 (default, Mar 15 2022, 12:22:08) \n[GCC 9.4.0]"
  },
  "vscode": {
   "interpreter": {
    "hash": "5b3ded1ccb95c1d9bd405e7b823d9e85424cde40fbb5985eb47e999ef50e15b4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
